#!/usr/bin/env python
# 给repeatmasker生成的重复序列重命名
from Bio import SeqIO
import argparse
import sys
parser = argparse.ArgumentParser(
    description='对repeatlib进行重新命名，由大天才于2020年5月3日创建于浙江农业大学')

parser.add_argument('-i',
                help='必须给定，输入的fasta文件， 其id为repeatmakser的xxx#class/family类型')
parser.add_argument('-p',
                help='必须给定，确定物种的前缀')
parser.add_argument('-d',
                help='默认家族名称，默认为repeat')
parser.add_argument('-o',
                help='必须给定，输出文件')



args = parser.parse_args()

if not args.i or not args.p or not args.o:
    parser.print_help()
    sys.exit()

infile = args.i
prefix = args.p

outfile = args.o

if not args.d:
    default_type = "repeat" 
else:
    default_type = args.d
dic = {}
with open(outfile,'w') as fila:
    for i in SeqIO.parse(infile,'fasta'):
        name = i.name
        k = name.split('#')
        if len(k)>1:
            the_type = k[1].split(' ')[0]
        else:
            the_type = default_type
        if the_type in dic:
            dic[the_type] +=1
        else:
            dic[the_type] = 1

        new_name = prefix+'.'+the_type+'.'+str(dic[the_type])
        fila.write('>'+new_name+'#'+the_type+'\n')
        fila.write(str(i.seq)+'\n')
